#include "nsbtree.h"

namespace nsbtree {

// local buffer for IO
thread_local static char local_buff[NSBTREE_PAGE_SIZE];

void Serialize(ReqType op, _key_t k, const string & v, char * buffer) {
    Request * req = (Request *)buffer;
    req->op = op; 
    req->crc = SIMPLE_CRC; 
    req->val_size = v.size();
    *((uint64_t *)req->keyvalue) = k;
    memcpy(req->keyvalue + 8, v.c_str(), v.size());
}

/*
    LeafPage interfaces
*/
_key_t LeafPage::Insert(const _key_t key, const string & val, LeafPage * new_leaf) {
    if(ShouldSplit(val)) {
        // split the node
        int m = count_ / 2;

        char tmp[NSBTREE_PAGE_SIZE] = {};
        LeafPage * this_copy = (LeafPage *)tmp;
        this_copy->Clear();
        this_copy->pageid_ = this->pageid_;
        this_copy->next_pageid_ = new_leaf->pageid_;
        this_copy->splitkey = recs_[m].key;
        for(int i = 0; i < m; i++) {
            std::string tmp(buff_ + recs_[i].val_off, recs_[i].val_size);
            auto discard = this_copy->Insert(recs_[i].key, tmp);
            assert(discard == 0);
        }

        for(int i = m; i < count_; i++) {
            std::string tmp(buff_ + recs_[i].val_off, recs_[i].val_size);
            auto discard = new_leaf->Insert(recs_[i].key, tmp);
            assert(discard == 0);
        }

        new_leaf->next_pageid_ = next_pageid_;
        new_leaf->splitkey = this->splitkey;
        memcpy(this, this_copy, sizeof(LeafPage) + this_copy->cur_off_);

        if(key < splitkey) {
            auto discard = Insert(key, val);
            assert(discard == 0);
        } else {
            auto discard = new_leaf->Insert(key, val);
            assert(discard == 0);
        }
        return this->splitkey;
    } else {
        int i = 0;
        for(; i < count_; i++) {
            if(recs_[i].key >= key) {
                break;
            }
        }

        if(recs_[i].key == key) {
            // udpate
            if(recs_[i].val_size >= val.size()) {
                // in place update
                memcpy(buff_ + recs_[i].val_off, val.c_str(), recs_[i].val_size);
            } else if(cur_off_ + val.size() <= MAX_OFF) {
                // contains enough space
                recs_[i].val_off = cur_off_;
                recs_[i].val_size = val.size();
                memcpy(buff_ + cur_off_, val.c_str(), val.size());
                cur_off_ += val.size();
            } else {
                // TODO: cannot update this value
            }
            return 0;
        } else {
            // insert
            memmove(recs_ + i + 1, recs_ + i, sizeof(Record) * (count_ - i));
            recs_[i].key = key;
            recs_[i].val_off = cur_off_;
            recs_[i].val_size = val.size();
            memcpy(buff_ + cur_off_, val.c_str(), val.size());
            cur_off_ += val.size();
            count_ += 1;
            return 0;
        }
    }
}

bool LeafPage::Lookup(const _key_t key, string & val) {
    int i = 0;
    for(; i < count_; i++) {
        if(recs_[i].key >= key) {
            break;
        }
    }

    if(recs_[i].key == key) {
        val = std::string(buff_ + recs_[i].val_off, recs_[i].val_size);
        return true;
    } else {
        return false;
    }
}

bool LeafPage::Update(const _key_t key, const string & val) {
    int i = 0;
    for(; i < count_; i++) {
        if(recs_[i].key >= key) {
            break;
        }
    }

    if(recs_[i].key == key) { // udpate
        if(recs_[i].val_size >= val.size()) {
            // in place update
            memcpy(buff_ + recs_[i].val_off, val.c_str(), recs_[i].val_size);
        } else if(cur_off_ + val.size() <= MAX_OFF) {
            // contains enough space
            recs_[i].val_off = cur_off_;
            recs_[i].val_size = val.size();
            memcpy(buff_ + cur_off_, val.c_str(), val.size());
            cur_off_ += val.size();
        } else {
            // TODO: cannot update this value
        }
        return true;
    } else {
        return false;
    }
}

int LeafPage::Scan(const _key_t startKey, int len, vector<std::string> & out) {
    int num = 0;
    for(int i = 0; i < count_ && out.size() < len; i++) {
        if(recs_[i].key >= startKey) {
            std::string tmp(buff_ + recs_[i].val_off, recs_[i].val_size);
            out.push_back(tmp);
            num += 1;
        }
    }
    return num;
}

void LeafPage::Compact() {
    char tmp[NSBTREE_PAGE_SIZE] = {};
    LeafPage * this_copy = (LeafPage *)tmp;
    this_copy->Clear();
    this_copy->pageid_ = pageid_;
    this_copy->next_pageid_ = next_pageid_;
    this_copy->splitkey = splitkey;
    for(int i = 0; i < count_; i++) {
        std::string tmp(buff_ + recs_[i].val_off, recs_[i].val_size);
        this_copy->Insert(recs_[i].key, tmp);
    }
    memcpy(this, this_copy, sizeof(LeafPage) + this_copy->cur_off_);
}

/*
    NSBtree interfaces
*/
NSBtree::NSBtree(NSBOpt & opt) {
    bool createIf = false;
    if(!folder_exist(opt.dbdir.c_str())) {
        mkdir(opt.dbdir.c_str(), 0777);
    }
    if(!folder_exist(opt.nvmdir.c_str())) {
        createIf = true;
        mkdir(opt.nvmdir.c_str(), 0777);
    }

    std::string dbname = opt.dbdir + "/nsbtree.db";
    db_fd_ = open(dbname.c_str(), O_RDWR | O_CREAT, 0755);
    if(db_fd_ < 3) {
        fprintf(stderr, "open log file\n");
        exit(-1);
    }
    cur_page_id_.store(file_size(dbname.c_str()) / NSBTREE_PAGE_SIZE);

    uptree_ = new tlbtree_up::TLBtree(opt.nvmdir + "/tlbtree.pool", opt.nvmsize);
    LOG(INFO) << opt.ssdsize / NSBTREE_PAGE_SIZE;
    locktable_ = new RWLock[opt.ssdsize / NSBTREE_PAGE_SIZE];
    if(createIf) {
        uint32_t new_pid = NewPage();
        LeafPage * first_leaf = (LeafPage *)local_buff;
        first_leaf->Clear();
        first_leaf->splitkey = MAX_KEY;
        first_leaf->pageid_ = new_pid;
        first_leaf->next_pageid_ = INVALID_PAGE_ID;
        WritePage(new_pid, local_buff);
        uptree_->insert(MIN_KEY, first_leaf->pageid_);
    }

    // write-ahead log
    std::string logname = opt.dbdir + "/nsbtree.log";
    log_fd_ = open(logname.c_str(), O_RDWR | O_CREAT, 0755);
    if(log_fd_ < 3) {
        fprintf(stderr, "open log file\n");
        exit(-1);
    }
    if(file_size(logname.c_str()) != 0) {
        // TODO: Redo Log replay
        auto ret = ftruncate(log_fd_, 0);
        assert(ret >= 0);
    }
    log_offset_ = 0;
}

NSBtree::~NSBtree() {
    auto ret = fsync(log_fd_);
    ret = ftruncate(log_fd_, 0);
    delete uptree_;
    delete locktable_;
}

void NSBtree::Put(const _key_t & key, const std::string & val, bool sync) {
    Serialize(PUT, key, val, local_buff);
    Append((Request *) local_buff, sync);
    uint32_t leaf_page_id = uptree_->lookup(key);

    put_retry:
    locktable_[leaf_page_id].WLock();
        ReadPage(leaf_page_id, local_buff);
        LeafPage * leaf = (LeafPage *)local_buff;
        if(key >= leaf->splitkey) {
            uint32_t old = leaf_page_id;
            leaf_page_id = leaf->next_pageid_;
            locktable_[old].UnWLock();
            goto put_retry;
        }

        if(leaf->ShouldSplit(val)) {
            uint32_t new_pid = NewPage();
            locktable_[new_pid].WLock();
                char new_buff[NSBTREE_PAGE_SIZE];
                LeafPage * new_leaf = (LeafPage *)new_buff;
                new_leaf->Clear();
                new_leaf->pageid_ = new_pid;
                _key_t splitkey = leaf->Insert(key, val, new_leaf);
                uptree_->insert(splitkey, new_leaf->pageid_);
                WritePage(new_pid, new_buff);
            locktable_[new_pid].UnWLock();
        } else {
            LeafPage * new_leaf = nullptr;
            _key_t discard = leaf->Insert(key, val, new_leaf);
        }
    WritePage(leaf_page_id, local_buff);
    locktable_[leaf_page_id].UnWLock();
}

bool NSBtree::Get(const _key_t & key, std::string & val) {
    uint32_t leaf_page_id = uptree_->lookup(key);
    
    get_retry:
    locktable_[leaf_page_id].RLock();
        ReadPage(leaf_page_id, local_buff);
        LeafPage * leaf = (LeafPage *)local_buff;
        if(key >= leaf->splitkey) {
            uint32_t old = leaf_page_id;
            leaf_page_id = leaf->next_pageid_;
            locktable_[old].UnRLock();
            goto get_retry;
        }
        bool foundIf = leaf->Lookup(key, val);
    locktable_[leaf_page_id].UnRLock();
    return foundIf;
}

bool NSBtree::Update(const _key_t & key, std::string & val, bool sync) {
    Serialize(PUT, key, val, local_buff);
    Append((Request *) local_buff, sync);
    uint32_t leaf_page_id = uptree_->lookup(key);
    
    update_retry:
    locktable_[leaf_page_id].WLock();
        ReadPage(leaf_page_id, local_buff);
        LeafPage * leaf = (LeafPage *)local_buff;
        if(key >= leaf->splitkey) {
            uint32_t old = leaf_page_id;
            leaf_page_id = leaf->next_pageid_;
            locktable_[old].UnWLock();
            goto update_retry;
        }
        bool foundIf = leaf->Update(key, val);
    if(foundIf)
        WritePage(leaf_page_id, local_buff);
    locktable_[leaf_page_id].UnWLock();
    return foundIf;
}

bool NSBtree::Delete(const _key_t & key, bool sync) {
    // TODO: not implemented
    return false;
}

int NSBtree::Scan(const _key_t startKey, int len, vector<std::string> & out) {
    uint32_t leaf_page_id = uptree_->lookup(startKey);

    int total = 0;
    locktable_[leaf_page_id].RLock();
    ReadPage(leaf_page_id, local_buff);
    while(true) {
        LeafPage * leaf = (LeafPage *) local_buff;
        total += leaf->Scan(startKey, len, out);
        if(total >= len) {
            locktable_[leaf_page_id].UnRLock();
            break;
        }
        uint32_t old = leaf_page_id;
        leaf_page_id = leaf->next_pageid_;
        locktable_[old].UnRLock();

        // go check next leaf node
        if(leaf_page_id == INVALID_PAGE_ID) {
            break;
        } else {
            locktable_[leaf_page_id].RLock();
            ReadPage(leaf_page_id, local_buff);
        }
    }
    return total;
}

void NSBtree::Append(Request * req, bool sync) {
    uint64_t addr = __sync_fetch_and_add(&log_offset_, req->Length());
    auto discard = pwrite(log_fd_, req, req->Length(), addr);
    if(sync) fsync(log_fd_);
}

void NSBtree::ReadPage(uint32_t page_id, char * buff) {
    auto discard = pread(db_fd_, buff, NSBTREE_PAGE_SIZE, (uint64_t)page_id * NSBTREE_PAGE_SIZE);
    assert(discard == NSBTREE_PAGE_SIZE);
}

void NSBtree::WritePage(uint32_t page_id, char * buff) {
    auto discard = pwrite(db_fd_, buff, NSBTREE_PAGE_SIZE, (uint64_t)page_id * NSBTREE_PAGE_SIZE);
    assert(discard == NSBTREE_PAGE_SIZE);
}

} // namspace NSBtree